package com.mjk.common.util;

import com.plumelog.core.TraceId;
import org.apache.dubbo.rpc.RpcContext;
import org.slf4j.MDC;

import java.util.UUID;
/**
 * MDC 工具类
 * @author mjk
 *
 */
public class MDCUtils {
	
	public final static String TRACE_ID = "trace_id";
	 
    /**
     * 初始化traceId,由consumer调用
     */
    public static void initMDC() {
        String traceId = generateTraceId();
        setTraceId(traceId);
    }
 
    /**
     * 从Dubbo中获取traceId，provider调用
     */
    public static void getMDC(RpcContext context) {
        String traceId = context.getAttachment(TRACE_ID);
        if (traceId == null) {
            traceId = generateTraceId();
        }
        setTraceId(traceId);
    }
 
    /**
     * 把traceId放入dubbo远程调用中，consumer调用
     */
    public static void putMDC(RpcContext context) {
        String traceId = getTraceId();
        if (traceId != null) {
            context.setAttachment(TRACE_ID, traceId);
        }
    }
 
    /**
     * 从MDC中清除traceId
     */
    public static void clearMDC() {
        MDC.remove(TRACE_ID);
    }

    /**
     * 从MDC中获取traceId
     * */
    public static String getTraceId() {
        return MDC.get(TRACE_ID);
    }
 
    /**
     * 将traceId放入MDC
     */
    public static void setTraceId(String traceId) {
        MDC.put(TRACE_ID, traceId);
        TraceId.logTraceID.set(traceId);
    }
 
    /**
     * 生成traceId
     */
    static private String generateTraceId() {
        return UUID.randomUUID().toString().replace("-", "");
    }

}
